package net.kehuilai.tools;

import com.amazon.sqs.javamessaging.AmazonSQSMessagingClientWrapper;
import com.amazon.sqs.javamessaging.SQSConnection;
import com.amazon.sqs.javamessaging.SQSConnectionFactory;
import com.amazon.sqs.javamessaging.SQSSession;
import com.amazonaws.auth.AWSCredentials;
import com.amazonaws.auth.AWSStaticCredentialsProvider;
import com.amazonaws.auth.policy.Policy;
import com.amazonaws.auth.policy.Principal;
import com.amazonaws.auth.policy.Resource;
import com.amazonaws.auth.policy.Statement;
import com.amazonaws.auth.policy.actions.SNSActions;
import com.amazonaws.auth.policy.actions.SQSActions;
import com.amazonaws.auth.profile.ProfileCredentialsProvider;
import com.amazonaws.regions.Region;
import com.amazonaws.regions.Regions;
import com.amazonaws.services.sns.AmazonSNS;
import com.amazonaws.services.sns.AmazonSNSClient;
import com.amazonaws.services.sns.model.SubscribeResult;
import com.amazonaws.services.sns.model.Topic;
import com.amazonaws.services.sqs.AmazonSQS;
import com.amazonaws.services.sqs.AmazonSQSClient;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.jms.Message;
import java.io.IOException;
import java.lang.reflect.Method;
import java.net.Inet4Address;
import java.net.InetAddress;
import java.net.NetworkInterface;
import java.net.SocketException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * Created by genghz on 16/11/11.
 */
public class AwsMqReceiver {

    private static final Logger log = LoggerFactory.getLogger(AwsMqReceiver.class);
    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
    static {
        OBJECT_MAPPER.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
    }

    private static final String uuid = UUID.randomUUID().toString();
    private static final String env = System.getenv().get("ENV");
    private static final AWSCredentials credentials = new ProfileCredentialsProvider().getCredentials();

    private SQSConnection connection;
    private AmazonSQS sqsClient;
    private AmazonSNS snsClient;

    private ConcurrentMap<String, String> queue2url = new ConcurrentHashMap<String, String>();

    private Set<String> queueDelete = new HashSet<String>();
    private Set<String> subIdDelete = new HashSet<String>();

    private List<String> topicArns = new ArrayList<String>();


    public AwsMqReceiver() {
        this(Regions.CN_NORTH_1);
    }

    public AwsMqReceiver(Regions regions) {
        try {

            Region region = Region.getRegion(regions);

            SQSConnectionFactory connectionFactory = SQSConnectionFactory.builder()
                    .withRegion(region)
                    .withAWSCredentialsProvider(new AWSStaticCredentialsProvider(credentials))
                    .build();
            this.connection = connectionFactory.createConnection();

            this.sqsClient = new AmazonSQSClient(credentials);
            sqsClient.setRegion(region);

            this.snsClient = new AmazonSNSClient(credentials);
            snsClient.setRegion(region);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    // topic <-> queue 一对一处理 同一时间 集群中只有一台服务器收到
    public void receive(String name, int receiverCount, ReceiverHandler receiver) {
        receive(name, receiverCount, receiver, "sqsMessage");
    }

    public void receive(String name, int receiverCount, ReceiverHandler receiver, String method) {
        try {
            receiverCount = receiverCount <= 0 ? 1 : receiverCount;
            String queueName = getName(name);
            // create queue
            createdQueue(queueName);
            // create topic
            String topicArn = createdTopic(queueName);
            // sub
            queueSubTopic(queueName, topicArn, false);

            // 创建 receiver
            for (int i = 0; i < receiverCount; i++) {
                Session session = this.connection.createSession(false, SQSSession.UNORDERED_ACKNOWLEDGE);
                Queue queue = session.createQueue(queueName);
                MessageConsumer consumer = session.createConsumer(queue);
                consumer.setMessageListener(new MessageListenerImpl(receiver, method, queueName));
            }

            log.info("init receiver: queue: {}, receiver_count: {}", queueName, receiverCount);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    // topic <-> queue 一对多订阅 同一时间 集群中每台服务器都收到
    public void broadcast(String name, int receiverCount, ReceiverHandler receiver) {
        broadcast(name, true, receiverCount, receiver, "sqsMessage");
    }

    public void broadcast(String name, int receiverCount, ReceiverHandler receiver, String method) {
        broadcast(name, true, receiverCount, receiver, method);
    }

    private void broadcast(String name, boolean autoDelete, int receiverCount, ReceiverHandler receiver, String method) {
        try {
            receiverCount = receiverCount <= 0 ? 1 : receiverCount;
            String topicName = getName(name);
            String queueName = getUUName(name);
            // create queue
            createdQueue(queueName);
            // create topic
            String topicArn = createdTopic(topicName);
            // sub
            queueSubTopic(queueName, topicArn, autoDelete);

            for (int i = 0; i < receiverCount; i++) {
                Session session = this.connection.createSession(false, SQSSession.UNORDERED_ACKNOWLEDGE);
                Queue queue = session.createQueue(queueName);
                MessageConsumer consumer = session.createConsumer(queue);
                consumer.setMessageListener(new MessageListenerImpl(receiver, method, queueName));
            }

            log.info("init receiver: queue: {}, receiver_count: {}", queueName, receiverCount);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public void start() {
        try {
            this.connection.start();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public void destroy() {
        for (String queueName : queueDelete) {
            this.sqsClient.deleteQueue(getQueueUrl(queueName));
        }

        for (String subId : subIdDelete) {
            this.snsClient.unsubscribe(subId);
        }
    }

    private void pushQueue(String queueName, String text, String classType, int delaySeconds) {
        String queueUrl = getQueueUrl(queueName);

        SendMessageRequest request = new SendMessageRequest(queueUrl, text);
        request.addMessageAttributesEntry("_type_", new MessageAttributeValue().withDataType("String").withStringValue(classType));
        request.setDelaySeconds(delaySeconds);
        sqsClient.sendMessage(request);
    }

    private void createdQueue(String name) {
        try {
            AmazonSQSMessagingClientWrapper client = this.connection.getWrappedAmazonSQSClient();
            if (!client.queueExists(name)) {
                CreateQueueResult result = client.createQueue(name);
                String queueUrl = result.getQueueUrl();
                this.queue2url.put(name, queueUrl);

                this.setSqsPolicy(queueUrl);
            }
        } catch (JMSException e) {
            e.printStackTrace();
        }
    }

    private String createdTopic(String name) {

        // 获取 arn
        if (this.topicArns.isEmpty()) {
            List<Topic> topics = this.snsClient.listTopics().getTopics();
            for (Topic topic : topics) {
                this.topicArns.add(topic.getTopicArn());
            }
        }

        // 创建获取 arn
        String topicArn = this.snsClient.createTopic(name).getTopicArn();
        if (this.topicArns.contains(topicArn)) return topicArn;

        // 设置sns 策略
        this.setSnsPolicy(topicArn);

        this.topicArns.add(topicArn);
        return topicArn;
    }

    private String getQueueUrl(String queueName) {
        if (!queue2url.containsKey(queueName)) {
            queue2url.put(queueName, this.sqsClient.getQueueUrl(queueName).getQueueUrl());
        }
        return this.queue2url.get(queueName);
    }

    // queue sub topic
    private void queueSubTopic(String queueName, String topicArn, boolean autoDelete) {
        List<String> attrs = new ArrayList<String>();
        attrs.add("QueueArn");
        String queueArn = this.sqsClient.getQueueAttributes(getQueueUrl(queueName), attrs).getAttributes().get("QueueArn");

        // 添加订阅
        SubscribeResult subResult = this.snsClient.subscribe(topicArn, "sqs", queueArn);

        // 设置原消息传输
        this.snsClient.setSubscriptionAttributes(subResult.getSubscriptionArn(), "RawMessageDelivery", "true");

        if (autoDelete) {
            this.queueDelete.add(queueName);
            this.subIdDelete.add(subResult.getSubscriptionArn());
        }
    }

    // 配置 sns 策略
    private void setSnsPolicy(String topicArn) {
        Policy snsPolicy =
                new Policy().withStatements(
                        new Statement(Statement.Effect.Allow)
                                .withPrincipals(Principal.AllUsers)
                                .withActions(SNSActions.Subscribe, SNSActions.Publish)
                                .withResources(new Resource(topicArn)));

        this.snsClient.setTopicAttributes(topicArn, "Policy", snsPolicy.toJson());
    }

    // 配置 sqs 策略
    private void setSqsPolicy(String queueUrl) {

        String arn = this.sqsClient.getQueueAttributes(
                new GetQueueAttributesRequest(queueUrl).withAttributeNames("QueueArn")
        ).getAttributes().get("QueueArn");

        Policy sqsPolicy =
                new Policy().withStatements(
                        new Statement(Statement.Effect.Allow)                           // 策略：允许
                                .withPrincipals(Principal.AllUsers)                     // 委托人(可执行人)：默认所有人
                                .withActions(SQSActions.AllSQSActions)                  // 委托人权限：默认全部权限
                                .withResources(new Resource(arn)));                     // 策略资源：队列arn

        Map<String, String> queueAttributes = new HashMap<String, String>();
        queueAttributes.put("Policy", sqsPolicy.toJson());

        this.sqsClient.setQueueAttributes(new SetQueueAttributesRequest(
                queueUrl,
                queueAttributes));
    }

    private String getName(String name) {
        if ("prod".equals(env)) {
            return name;
        } else {
            return name + "-test";
        }
    }

    private String getUUName(String name) {
        String name1 = getName(name);
        return name1 + "-" + uuid;
    }

    private void getIp() throws SocketException {
        Enumeration allNetInterfaces = NetworkInterface.getNetworkInterfaces();
        while (allNetInterfaces.hasMoreElements()) {
            NetworkInterface netInterface = (NetworkInterface) allNetInterfaces.nextElement();

            Enumeration addresses = netInterface.getInetAddresses();
            while (addresses.hasMoreElements()) {
                InetAddress ip = (InetAddress) addresses.nextElement();
                if (ip != null && ip instanceof Inet4Address && !ip.getHostAddress().equals("127.0.0.1")) {
                    System.out.println(ip.getHostAddress());
                }
            }
        }
    }

    private <T> T json2pojo(String json, Class<T> type) {
        if (json == null || "".equals(json)) return null;
        T t = null;
        try {
            json = json.trim();
            t = OBJECT_MAPPER.readValue(json, type);
        } catch (IOException e) {
            log.warn("json2pojo failed: {}", e.getMessage());
        }
        return t;
    }

    private class MessageListenerImpl implements MessageListener {

        private ReceiverHandler handler;
        private String method;
        private String queue;

        public MessageListenerImpl(ReceiverHandler handler, String method, String queue) {
            this.handler = handler;
            this.method = method;
            this.queue = queue;
        }

        public void onMessage(Message message) {

            try {
                if (message instanceof TextMessage) {
                    TextMessage txtMessage = (TextMessage) message;

                    String delayStr = message.getStringProperty("_delay_");
                    String classType = message.getStringProperty("_type_");
                    // 延时消息
                    if (delayStr != null && !delayStr.isEmpty()) {
                        int delaySeconds = Integer.valueOf(delayStr);
                        if (delaySeconds > 0) {
                            pushQueue(this.queue, txtMessage.getText(), classType, delaySeconds);
                            message.acknowledge();
                            log.info("receive delay message: {}s, {}", delayStr, txtMessage.getText());
                            return;
                        }
                    }
                    invoke(txtMessage);
                    message.acknowledge();
                } else {
                    throw new ClassNotFoundException("receive message not TextMessage");
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }

        private void invoke(TextMessage message) throws Exception {
            String classType = message.getStringProperty("_type_");
            Object o = null;
            Class T = null;
            if (classType != null && !classType.equals("")) {
                T = Class.forName(classType);
                o = json2pojo(message.getText(), T);
            } else {
                T = String.class;
                o = message.getText();
            }

            // 分发
            Method method = handler.getClass().getMethod(this.method, T);
            method.invoke(handler, o);
        }
    }

    public interface ReceiverHandler {
    }
}

